import sqlalchemy


def get_pre_lessons(lesson_num: str, engine):
    conn = engine.connect()
    try:
        result = conn.execute(sqlalchemy.text("""
                    WITH RECURSIVE SubsequentLessons AS (
                      -- 非递归成员（初始查询）
                      SELECT to_lesson
                      FROM lesson_prerequisite
                      WHERE from_lesson = :lesson_num
                      UNION ALL
                      -- 递归成员
                      SELECT lp.to_lesson
                      FROM lesson_prerequisite lp
                      INNER JOIN SubsequentLessons sl ON lp.from_lesson = sl.to_lesson
                    )
                    -- 最终查询，选择递归CTE的结果
                    SELECT DISTINCT * FROM SubsequentLessons ORDER BY to_lesson;
                                        """),
                              {"lesson_num": lesson_num})

        conn.commit()
        pre_lesson_list = result.fetchall()

        if len(pre_lesson_list) == 0:
            return None
        return pre_lesson_list

    except Exception as e:
        print(e)
        raise e
    finally:
        conn.close()


def get_pre_lessons_grades(lesson_id: str, engine):
    pre_lessons = get_pre_lessons(lesson_id, engine)
    if not pre_lessons:
        return {}

    grades_dict = {}
    for pre_lesson in pre_lessons:
        conn = engine.connect()
        try:
            query = f"""
                   SELECT sg.uid, sg.grade, sg.lesson_num, sg.semester
                   FROM stu_grade sg
                   WHERE sg.lesson_num = '{pre_lesson[0]}'
                   ORDER BY sg.uid,sg.semester;
               """
            result = conn.execute(sqlalchemy.text(query))

            for row in result.fetchall():
                student_id = str(row[0])
                grade = float(row[1])
                if student_id not in grades_dict:
                    grades_dict[student_id] = [None] * len(pre_lessons)
                index = pre_lessons.index(pre_lesson)
                grades_dict[student_id][index] = grade
        except Exception as e:
            print(e)
            raise e
        finally:
            conn.close()

    # 过滤掉那些没有完整成绩记录的学生
    grades_dict = {k: v for k, v in grades_dict.items() if None not in v}

    return grades_dict


def get_grade_for_stu(uid: str, lesson_num: str, engine):
    conn = engine.connect()
    try:
        query = sqlalchemy.text("""
            SELECT grade
            FROM stu_grade
            WHERE uid = :uid AND lesson_num = :lesson_num
            ORDER BY id DESC
            LIMIT 1;
        """)
        result = conn.execute(query, {"uid": uid, "lesson_num": lesson_num})

        row = result.fetchone()
        if row:
            return row[0]
        else:
            return None

    except Exception as e:
        print(e)
        raise e
    finally:
        conn.close()


def get_stu_in_class(class_name,engine):
    conn = engine.connect()
    try:
        query = sqlalchemy.text("""
                SELECT stu_detail.uid
                FROM stu_detail
                WHERE stu_detail.class_name = :class_name;
            """)
        result = conn.execute(query, {"class_name": class_name})
        uid_in_class = [row[0] for row in result.fetchall()]
        return  uid_in_class
    except Exception as e:
        print(e)
        raise e
    finally:
        conn.close()


def get_stu_info_dict(class_name, engine):
    conn = engine.connect()
    stu_list_result = conn.execute(sqlalchemy.text("""
                        SELECT stu_detail.uid, stu_detail.stu_num, stu_detail.major,
                         stu_detail.stu_position, user_info.username
                        FROM stu_detail, user_info
                        WHERE stu_detail.class_name = :class_name and stu_detail.uid = user_info.uid
                    """), {
        "class_name": class_name
    })
    conn.commit()
    stu_list = stu_list_result.fetchall()
    stu_info_dict = {stu[0]: {"stu_num": stu[1], "username": stu[4]} for stu in stu_list}
    return stu_info_dict


def generate_training_dataset(lesson_num: str, engine):
    pre_lesson_grades = get_pre_lessons_grades(lesson_num, engine)
    current_lesson_grades = {}
    for uid in pre_lesson_grades.keys():
        grade = get_grade_for_stu(uid, lesson_num, engine)
        if grade is not None:
            current_lesson_grades[uid] = grade

    x = []
    y = []
    for uid, grades in pre_lesson_grades.items():
        if uid in current_lesson_grades:
            x.append(grades)
            y.append(current_lesson_grades[uid])

    return x, y